package com.jojo.rateLimiter;

import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;

/**
 * 滑动窗口算法
 *
 * @author hpl
 * @since 2024/12/12
 */
public class SlidingWindow {
    private final int maxRequests;  // 每个时间窗口内最大请求数
    private final int bucketCount;  // 小桶数量
    private final long windowSize;  // 时间窗口大小（单位：毫秒）
    private final long bucketSize;  // 每个小桶的大小（单位：毫秒）

    // 每个小桶的请求计数器
    private final AtomicInteger[] requestCounts;
    private final Lock lock;  // 锁用于保证线程安全

    private long lastTimestamp;  // 上次更新时间戳

    public SlidingWindow(int maxRequests, long windowSize, int bucketCount) {
        this.maxRequests = maxRequests;
        this.windowSize = windowSize;
        this.bucketCount = bucketCount;
        this.bucketSize = windowSize / bucketCount;
        this.requestCounts = new AtomicInteger[bucketCount];
        this.lock = new ReentrantLock();

        // 初始化请求计数器
        for (int i = 0; i < bucketCount; i++) {
            requestCounts[i] = new AtomicInteger(0);
        }

        this.lastTimestamp = System.currentTimeMillis();
    }

    // 判断是否允许请求
    public boolean allowRequest() {
        lock.lock();
        try {
            long currentTime = System.currentTimeMillis();
            // 更新小桶状态
            updateBuckets(currentTime);

            // 计算整个窗口内的总请求数
            int totalRequests = 0;
            for (AtomicInteger count : requestCounts) {
                totalRequests += count.get();
            }

            if (totalRequests < maxRequests) {
                // 将请求放入当前小桶
                int currentBucket = (int) ((currentTime / bucketSize) % bucketCount);
                requestCounts[currentBucket].incrementAndGet();
                return true;  // 允许请求
            }
            return false;  // 超过请求限制
        } finally {
            lock.unlock();
        }
    }

    // 更新小桶，滑动窗口
    private void updateBuckets(long currentTime) {
        long timeElapsed = currentTime - lastTimestamp;

        // 如果有时间滑动，更新小桶
        if (timeElapsed > bucketSize) {
            int bucketsToUpdate = (int) (timeElapsed / bucketSize);
            for (int i = 0; i < bucketsToUpdate; i++) {
                int nextBucketIndex = (int) ((lastTimestamp / bucketSize + i) % bucketCount);
                requestCounts[nextBucketIndex].set(0);  // 清空最旧的小桶的请求计数
            }
            lastTimestamp = currentTime;  // 更新最后更新时间戳
        }
    }

    public static void main(String[] args) {
        int maxRequests = 5;  // 每秒最多允许 5 个请求
        long windowSize = 1000;  // 时间窗口为 1000 毫秒（1 秒）
        int bucketCount = 10;  // 将时间窗口划分为 10 个小桶

        SlidingWindow rateLimiter = new SlidingWindow(maxRequests, windowSize, bucketCount);

        // 模拟 20 次请求
        for (int i = 0; i < 20; i++) {
            if (rateLimiter.allowRequest()) {
                System.out.println("Request " + (i + 1) + " allowed");
            } else {
                System.out.println("Request " + (i + 1) + " denied");
            }
            try {
                Thread.sleep(100);  // 每 100 毫秒请求一次
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
        }
    }
}
